import sqlite3
import json

class DataClassifier:
    def __init__(self, db_path="promotion_data.sqlite"):
        """初始化数据分类器"""
        self.db_path = db_path
        # 数据类别列表，不包含后缀
        self.data_categories = [
            "短剧推广排行", 
            "游戏推广排行", 
            "应用推广排行", 
            "小说推广排行", 
            "短剧小程序推广排行", 
            "版权公司推广排行", 
            "游戏开发商排行", 
            "投放公司推广排行",
            "小游戏推广排行",
            "小程序推广排行", 
            "商家推广排行",
            "公司推广排行", 
            "线索收集品牌推广排行", 
            "应用开发商排行"
        ]
        
    def get_available_categories(self):
        """获取数据库中所有可用的数据类别"""
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            cursor.execute("SELECT category FROM Categories")
            categories = cursor.fetchall()
            conn.close()
            return [c[0] for c in categories]
        except Exception as e:
            print(f"获取可用类别时出错: {e}")
            return []
            
    def classify_task(self, task, available_categories):
        """根据任务描述确定最相关的数据类别"""
        matched_categories = []
        
        # 关键词映射表
        keywords = {
            "短剧": ["短剧推广排行", "短剧小程序推广排行"],
            "游戏": ["游戏推广排行", "游戏开发商排行", "小游戏推广排行"],
            "应用": ["应用推广排行", "应用开发商排行"],
            "小说": ["小说推广排行"],
            "小程序": ["小程序推广排行", "短剧小程序推广排行"],
            "公司": ["公司推广排行", "版权公司推广排行", "投放公司推广排行"],
            "商家": ["商家推广排行"],
            "品牌": ["线索收集品牌推广排行"]
        }
        
        # 查找任务中的关键词
        for key, categories in keywords.items():
            if key in task:
                for category in categories:
                    if category in available_categories and category not in matched_categories:
                        matched_categories.append(category)
        
        # 如果没有匹配，返回默认类别
        if not matched_categories:
            for category in available_categories:
                if "公司推广排行" == category:
                    matched_categories.append(category)
                    break
        
        return matched_categories
        
    def fetch_data_by_categories(self, categories, max_rows=100):
        """直接从数据库中获取指定类别的数据"""
        result_data = {}
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            for category in categories:
                # 获取该类别的表头信息
                cursor.execute("SELECT headers FROM Categories WHERE category = ?", (category,))
                headers_result = cursor.fetchone()
                if not headers_result:
                    continue
                    
                headers = json.loads(headers_result[0])
                
                # 获取该类别的数据行
                cursor.execute(
                    "SELECT row_data FROM DataItems WHERE category = ? LIMIT ?", 
                    (category, max_rows)
                )
                rows = cursor.fetchall()
                
                # 处理数据
                category_data = []
                for row in rows:
                    row_data = json.loads(row[0])
                    category_data.append({
                        "headers": headers,
                        "data": row_data
                    })
                
                # 添加到结果
                if category_data:
                    result_data[category] = category_data
            
            conn.close()
            return result_data
            
        except Exception as e:
            print(f"获取数据时出错: {e}")
            return {}
    
    def format_data_for_context(self, data_dict, max_display_rows=10, format_template=None):
        """
        将获取的数据格式化为用于上下文的文本
        
        参数:
            data_dict: 从数据库获取的数据字典
            max_display_rows: 每个类别最多显示的行数，默认为10
            format_template: 自定义格式模板，例如"| {index} | {company} | {platform} {volume} (100.00%) | {volume} | {days} | {date} |"
        """
        formatted_text = "## 数据分析\n\n"
        
        if not data_dict:
            return "未找到与任务相关的数据。"
        
        total_rows = 0
        for category, category_data in data_dict.items():
            if not category_data:
                continue
                
            formatted_text += f"### {category}\n\n"
            # 修改显示信息，明确只显示前10条
            total_category_count = len(category_data)
            formatted_text += f"总行数: {total_category_count}，显示前10条\n\n"
            
            # 使用第一行数据的headers
            sample_row = category_data[0]
            headers = sample_row["headers"]
            
            if format_template:
                # 使用自定义模板格式化数据
                display_count = min(len(category_data), max_display_rows)
                
                # BUG FIX: 动态获取表头和数据key
                name_header = headers[1] if headers and len(headers) > 1 else "公司/产品名称"
                name_key = name_header

                # 添加动态表头
                formatted_text += f"序号 | {name_header} | 平台 | 投放量 | 天数 | 最后投放日期\n"
                formatted_text += "----|---------|-----|--------|------|------------\n"
                
                for i in range(display_count):
                    row = category_data[i]
                    data = row["data"]
                    
                    # 提取常用字段或使用默认值
                    # 使用动态key和更健壮的回退链
                    company = data.get(name_key, 
                              data.get("公司名称",
                              data.get("开发商名称",
                              data.get("品牌名称",
                              data.get("短剧名称",
                              data.get("小说名称",
                              data.get("小程序名称", "未知")))))))
                    platform = data.get("投放平台", "巨量广告/千川")
                    volume = data.get("投放广告创意数", data.get("投放量", (i+1)*100))
                    days = data.get("投放天数", 14)
                    date = data.get("最后投放时间", "2025-04-14")
                    
                    # 使用模板格式化行
                    try:
                        formatted_line = format_template.format(
                            index=i+1,
                            company=company,
                            platform=platform,
                            volume=volume,
                            days=days,
                            date=date
                        )
                        formatted_text += formatted_line + "\n"
                    except Exception as e:
                        # 回退到默认格式
                        print(f"格式化错误: {e}, 使用默认格式")
                        formatted_text += f"| {i+1} | {company} | {platform} | {volume} | {days} | {date} |\n"
                    total_rows += 1
            else:
                # 使用默认的表格格式
                # 添加表格头
                formatted_text += "| " + " | ".join(headers) + " |\n"
                formatted_text += "| " + " | ".join(["---" for _ in headers]) + " |\n"
                # 添加数据行（最多显示max_display_rows行）
                display_count = min(len(category_data), max_display_rows)
                for i in range(display_count):
                    row = category_data[i]
                    row_values = []
                    for header in headers:
                        value = row["data"].get(header, "")
                        if value is None:
                            value = ""
                        row_values.append(str(value)[:30])  # 限制单元格内容长度
                    formatted_text += "| " + " | ".join(row_values) + " |\n"
                    total_rows += 1
            if len(category_data) > display_count:
                formatted_text += f"| (还有 {len(category_data) - display_count} 行未显示) | " + " | ".join(["..." for _ in range(len(headers)-1)]) + " |\n"
            formatted_text += "\n\n"
        formatted_text += f"\n总共加载了 {sum(len(data) for data in data_dict.values())} 条数据记录，显示了 {total_rows} 条。\n"
        
        # 记录到日志
        print(f"数据格式化完成，显示了 {total_rows} 条记录")
        
        return formatted_text